#include "crypto_tool.h"
#include "mpz_tool.h"

string CryptoTool::Hash(const char * input, size_t len)
{
  unsigned char * pch = new unsigned char[SHA_DIGEST_LENGTH];
  SHA1((const unsigned char*)input, len, pch);
  string output((char *)pch, (size_t)SHA_DIGEST_LENGTH);
  delete[] pch;
  return output;
}


string CryptoTool::Hash(const string& input)
{
  unsigned char * pch = new unsigned char[SHA_DIGEST_LENGTH];
  SHA1((const unsigned char*)input.c_str(), input.size(), pch);
  string output((char *)pch, (size_t)SHA_DIGEST_LENGTH);
  delete[] pch;
  return output;
}

string CryptoTool::Hash(const mpz_t * input)
{
  size_t bits = mpz_sizeinbase (*input, 2);
  size_t len = (bits+7)/8;
  char * input_buf = new char[len];
  mpz_export((void *)input_buf, NULL, 1, 1, 0, 0, *input);
  unsigned char * pch = new unsigned char[SHA_DIGEST_LENGTH];
  SHA1((const unsigned char*)input_buf, len, pch);
  string output((char *)pch, 160/8);
  delete[] input_buf;
  delete[] pch;
  return output;
}

string CryptoTool::FDH(const string& input, mpz_t * modulus)
{
  int nbits = mpz_sizeinbase (*modulus, 2);
  int repetition = (nbits + HASH_BITS_LEN - 1) / HASH_BITS_LEN;
  stringstream result;
  for (int i = 0; i < repetition; i++)
    {
	  string str = input + string((const char *)&i, sizeof(int));
      result<<Hash(str);
    }
  mpz_t * tmp = (mpz_t * ) malloc(sizeof(mpz_t));
  mpz_init(*tmp);
  MpzTool::Mpz_from_string(tmp, result.str());
  mpz_mod(*tmp, *tmp, *modulus);
  string rt;
  MpzTool::String_from_mpz(&rt, tmp);
  mpz_clear(*tmp);
  free(tmp);
  return rt;
}

void CryptoTool::FDH(mpz_t * mpzhash, const string& input, mpz_t * modulus)
{
  int nbits = mpz_sizeinbase (*modulus, 2);
  int repetition = (nbits + HASH_BITS_LEN - 1) / HASH_BITS_LEN;
  stringstream result;
  for (int i = 0; i < repetition; i++)
    {
	  string str = input + string((const char *)&i, sizeof(int));
      result<<Hash(str);
    }
  MpzTool::Mpz_from_string(mpzhash, result.str());
  mpz_mod(*mpzhash, *mpzhash, *modulus);
}

string CryptoTool::HTag(const string& input)
{
  stringstream ss;
  ss<<"1"<<input;
  return Hash(ss.str());
}

string CryptoTool::HKey(const string& input)
{
  stringstream ss;
  ss<<"2"<<input;
  return Hash(ss.str());
}

string CryptoTool::HIndKey(const string& input)
{
  stringstream ss;
  ss<<"3"<<input;
  return Hash(ss.str());
}

string CryptoTool::AES_CBC_Encrypt(const string& input, const string& key)
{
  unsigned long datalen = input.size();
  size_t modular = datalen % (MAX_BLOCK_SIZE);
  if ( modular != 0)
	datalen += (MAX_BLOCK_SIZE) - modular; // guarantee multiple of blocksize(16)
  char * output = new char[datalen];
  AES_KEY aes_key;
  unsigned char * ivec = new unsigned char [2*MAX_BLOCK_SIZE]; // default to 0
  memset(ivec, 0, 2*MAX_BLOCK_SIZE);
  AES_set_encrypt_key((const unsigned char *)key.c_str(), 128, &aes_key);
  AES_cbc_encrypt((const unsigned char *)input.c_str(), (unsigned char *)output, datalen, &aes_key, ivec, AES_ENCRYPT);
  delete [] ivec;
  string str(output, datalen);
  delete[] output;
  return str;
}

string CryptoTool::AES_CBC_Decrypt(const string& input, const string& key)
{
  unsigned long datalen = input.size();
  assert( datalen % (MAX_BLOCK_SIZE) == 0);
  char * output = new char[datalen];
  AES_KEY aes_key;
  unsigned char * ivec = new unsigned char [2*MAX_BLOCK_SIZE]; // default to 0
  memset(ivec, 0, 2*MAX_BLOCK_SIZE);
  AES_set_decrypt_key((const unsigned char *)key.c_str(), 128, &aes_key);
  AES_cbc_encrypt((const unsigned char *)input.c_str(), (unsigned char *)output, datalen, &aes_key, ivec, AES_DECRYPT);
  delete [] ivec;
  string str(output, datalen);	// up to application to do trimming
  delete[] output;
  return str;
}

string CryptoTool::AES_ECB_Encrypt(const string& input, const string& key)
{
  unsigned long datalen = input.size();
  char * output = new char[datalen];
  AES_KEY aes_key;
  AES_set_encrypt_key((const unsigned char *)key.c_str(), 128, &aes_key);
  AES_ecb_encrypt((const unsigned char *)input.c_str(), (unsigned char *)output, &aes_key, AES_ENCRYPT);
  string str(output, datalen);
  delete[] output;
  return str;
}

string CryptoTool::AES_ECB_Decrypt(const string& input, const string& key)
{
  unsigned long datalen = input.size();
  char * output = new char[datalen];
  AES_KEY aes_key;
  AES_set_decrypt_key((const unsigned char *)key.c_str(), 128, &aes_key);
  AES_ecb_encrypt((const unsigned char *)input.c_str(), (unsigned char *)output, &aes_key, AES_DECRYPT);
  string str(output, datalen);
  delete[] output;
  return str;  
}

string CryptoTool::RC4(const string& input, const string& key)
{
  unsigned long datalen = input.length();
  char * output = new char[datalen];	// get the same length as input
  RC4_KEY rc_key;
  RC4_set_key(&rc_key, key.length(), (const unsigned char *)key.c_str());
  ::RC4(&rc_key, datalen, (const unsigned char *)input.c_str(), (unsigned char *)output);
  string str(output, datalen);
  delete[] output;
  return str;
}


void CryptoTool::RC4(char * output, const char* input, const size_t datalen, const string& key)
{
  RC4_KEY rc_key;
  RC4_set_key(&rc_key, key.length(), (const unsigned char *)key.c_str());
  ::RC4(&rc_key, datalen, (const unsigned char *)input, (unsigned char *)output);
}


string CryptoTool::RC4(const string& input, const mpz_t * key)
{
  size_t bits = mpz_sizeinbase (*key, 2);
  size_t len = (bits+7)/8;
  char * buf = new char[len];
  size_t reallen = 0;
  mpz_export((void *)buf, &reallen, 1, 1, 0, 0, *key);
  assert(len == reallen);

  unsigned long datalen = input.length();
  char * output = new char[datalen];	// get the same length as input
  RC4_KEY rc_key;
  RC4_set_key(&rc_key, len, (const unsigned char *)buf);
  ::RC4(&rc_key, datalen, (const unsigned char *)input.c_str(), (unsigned char *)output);
  string str(output, datalen);
  delete[] output;
  delete[] buf;
  return str;
}
// assume input and key is 128-bit
string CryptoTool::XOR(const string& input, const string& key)
{
  size_t sz = input.size();
  char * output = new char[sz];
  long long * dout = (long long *)output;
  long long * din = (long long *)input.c_str();
  long long * dkey = (long long *)key.c_str();

  for (int i = 0; i < 2; i++)
	{
	  *dout = (*din) ^ (*dkey);
	  dout++;
	  din++;
	  dkey++;
	}
  string ret(output, sz);
  delete[] output;
  return ret;
}
